#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import yaml
import subprocess
import os
import sys
import signal # <--- 需要导入 signal 模块
from std_srvs.srv import Trigger, TriggerResponse
import atexit # <--- 也可以考虑使用 atexit 作为备用

class SingleServiceRunner:
    def __init__(self):
        # --- 省略了之前的初始化代码 (获取服务名, 加载配置等) ---
        # 1. Get the target service name from command-line arguments
        if len(sys.argv) < 2:
            print("错误：缺少必要的参数：target_service_name")
            print("用法：rosrun service_launcher single_service_runner.py <target_service_name> [_config_file:=/path/to/your/config.yaml]")
            sys.exit(1) # Exit if argument is missing
        self.target_service_name = sys.argv[1]

        # Initialize ROS node
        node_name = f"service_runner_{self.target_service_name}"
        node_name = node_name.replace('-', '_').replace('/', '_') # Sanitize node name
        # 使用 anonymous=True 确保即使意外运行多个同名 runner 也能启动，但服务名仍需唯一
        rospy.init_node(node_name, anonymous=True, log_level=rospy.INFO)
        rospy.loginfo(f"正在初始化服务运行器，目标服务: {self.target_service_name}")


        # 2. Get YAML configuration file path
        default_yaml_path = "/home/tl/ZD/ros_ws/src/service_launcher/config/config.yaml" # 使用你指定的默认路径
        self.yaml_path = rospy.get_param('~config_file', default_yaml_path)
        rospy.loginfo(f"将从以下路径加载服务配置: {self.yaml_path}")

        # 3. Load YAML and find the specific service configuration
        self.service_config = None
        try:
            # --- 省略了加载和验证 YAML 的代码 ---
            with open(self.yaml_path, 'r') as f:
                config = yaml.safe_load(f)
                if not config or 'services' not in config:
                    rospy.logerr(f"配置文件 {self.yaml_path} 格式无效。缺少 'services' 键。")
                    sys.exit(1)

                for cfg in config['services']:
                    if cfg.get('name') == self.target_service_name:
                        self.service_config = cfg
                        break

                if self.service_config is None:
                    rospy.logerr(f"在配置文件 {self.yaml_path} 中未找到服务名 '{self.target_service_name}'")
                    sys.exit(1)

                # Validate service config based on type
                service_type = self.service_config.get('type')
                if not service_type:
                     rospy.logerr(f"服务 '{self.target_service_name}' 的配置无效。缺少 'type' 键。")
                     sys.exit(1)
                # ... (省略其他类型检查) ...
                elif service_type == 'python':
                     if 'script_path' not in self.service_config:
                        rospy.logerr(f"python 服务 '{self.target_service_name}' 缺少 'script_path' 键")
                        sys.exit(1)
                elif service_type == 'bash':
                     if 'script_path' not in self.service_config:
                        rospy.logerr(f"bash 服务 '{self.target_service_name}' 缺少 'script_path' 键")
                        sys.exit(1)
                # ...

        except Exception as e: # 简化异常处理块
            rospy.logerr(f"初始化过程中发生错误: {e}")
            sys.exit(1)

        rospy.loginfo(f"已找到 '{self.target_service_name}' 的配置: {self.service_config}")

        # --- 新增：用于存储此实例启动的子进程 ---
        self.launched_processes = []
        # ---------------------------------------

        # 4. 创建指定的服务
        service_name_ros = f"/{self.target_service_name}"
        try:
            # 检查服务是否已存在 (避免因 anonymous=True 导致多个 runner 监听同一个服务)
            # 注意：这在高并发下可能有竞态条件，但对于手动启动基本足够
            try:
                rospy.wait_for_service(service_name_ros, timeout=0.1)
                rospy.logerr(f"错误：服务 '{service_name_ros}' 似乎已由另一个节点提供。请确保每个服务只有一个监听器。")
                sys.exit(1)
            except rospy.ROSException:
                # 服务不存在，是好事，可以继续创建
                pass

            rospy.Service(service_name_ros, Trigger, self.handle_request)
            rospy.loginfo(f"已成功创建并监听服务: {service_name_ros}")
        except Exception as e:
             rospy.logerr(f"创建服务 '{service_name_ros}' 失败: {e}")
             sys.exit(1)

        # --- 新增：注册关闭钩子 ---
        rospy.on_shutdown(self.shutdown_hook)
        # 也可以考虑使用 atexit，它在 Python 解释器退出时执行，作为备用方案
        # atexit.register(self.shutdown_hook)
        # --------------------------

    def handle_request(self, req):
        # --- 省略了构建命令的代码 ---
        service_name = self.service_config['name']
        service_type = self.service_config['type']
        params = self.service_config.get('params', []) # 用于 rosrun/roslaunch
        args = self.service_config.get('args', [])     # 用于 python
        working_dir = self.service_config.get('working_dir', None)
        command = []
        # ... (构建 command 列表的代码) ...
        if service_type == "rosrun":
            package = self.service_config['package']
            node = self.service_config['node']
            command = ['rosrun', package, node] + params
        elif service_type == "roslaunch":
            package = self.service_config['package']
            launch_file = self.service_config['launch_file']
            command = ['roslaunch', package, launch_file] + params
        elif service_type == "python":
            script_path = self.service_config['script_path']
            command = ['python', script_path] + args
        elif service_type == "bash":
            script_path = self.service_config['script_path']
            command = ['bash', script_path] + args
        else:
            rospy.logerr(f"内部错误：服务 '{service_name}' 的类型 '{service_type}' 不支持")
            return TriggerResponse(success=False, message=f"内部错误：不支持的类型 {service_type}")


        # 执行命令
        try:
            rospy.loginfo(f"正在为服务 '{service_name}' (类型: {service_type}) 执行命令: {' '.join(command)}")
            effective_cwd = None
            # ... (设置 effective_cwd 的代码) ...
            if working_dir:
                if os.path.isdir(working_dir):
                    rospy.loginfo(f"设置工作目录为: {working_dir}")
                    effective_cwd = working_dir
                else:
                     rospy.logwarn(f"工作目录 '{working_dir}' 不存在或不是一个目录。将在默认目录运行命令。")

            # 使用 preexec_fn=os.setsid 使得子进程成为新的进程组领头进程
            # 这对于后续使用 os.killpg 关闭整个进程组（特别是 roslaunch 的子孙进程）很有用
            process = subprocess.Popen(command, cwd=effective_cwd, preexec_fn=os.setsid)
            pid = process.pid
            rospy.loginfo(f"已成功为 '{service_name}' 启动命令，进程 PID: {pid} (进程组 PGID: {os.getpgid(pid)})")

            # --- 新增：将新启动的进程对象添加到列表中 ---
            self.launched_processes.append(process)
            # ------------------------------------------

            # --- 新增：定期清理已结束的进程 (可选但推荐) ---
            self.cleanup_finished_processes()
            # ---------------------------------------------

            return TriggerResponse(success=True, message=f"已启动 {service_name} (PID: {pid})")

        except Exception as e: # 简化异常处理
            rospy.logerr(f"为服务 '{service_name}' 启动命令时失败: {e}")
            # 确保在 FileNotFoundError 等情况下也能捕获
            return TriggerResponse(success=False, message=f"启动 {service_name} 时出错: {e}")

    def cleanup_finished_processes(self):
        """移除列表中已经结束的进程"""
        # 创建一个新列表只包含仍在运行的进程
        running_processes = []
        for proc in self.launched_processes:
            if proc.poll() is None: # poll() 返回 None 表示仍在运行
                running_processes.append(proc)
            else:
                rospy.loginfo(f"进程 PID {proc.pid} 已结束 (返回码: {proc.returncode})，将其从监控列表移除。")
        self.launched_processes = running_processes

    def shutdown_hook(self):
        """在节点关闭时被调用，用于清理所有由此实例启动的子进程"""
        rospy.loginfo(f"服务运行器 '{self.target_service_name}' 正在关闭，开始清理子进程...")

        # 迭代列表的一个副本，因为我们可能会在循环中修改列表（虽然当前实现不需要）
        for process in list(self.launched_processes):
            # 再次检查进程是否还在运行，以防在调用 shutdown_hook 前它刚好结束
            if process.poll() is None:
                pid = process.pid
                try:
                    pgid = os.getpgid(pid) # 获取进程组 ID
                    rospy.loginfo(f"正在尝试终止进程组 PGID: {pgid} (领头进程 PID: {pid})")

                    # 优先尝试发送 SIGINT (Ctrl+C) 给整个进程组，允许优雅关闭
                    os.killpg(pgid, signal.SIGINT)
                    rospy.loginfo(f"已发送 SIGINT 到 PGID {pgid}")

                    # 等待一小段时间看是否能正常退出
                    try:
                        process.wait(timeout=2.0) # 等待最多 2 秒
                        rospy.loginfo(f"进程组 PGID {pgid} (PID {pid}) 已成功终止。")
                    except subprocess.TimeoutExpired:
                        rospy.logwarn(f"进程组 PGID {pgid} (PID {pid}) 在 SIGINT 后 2 秒内未终止。尝试发送 SIGTERM...")
                        # 如果 SIGINT 不行，尝试发送 SIGTERM
                        os.killpg(pgid, signal.SIGTERM)
                        try:
                            process.wait(timeout=1.0) # 再等 1 秒
                            rospy.loginfo(f"进程组 PGID {pgid} (PID {pid}) 在 SIGTERM 后已终止。")
                        except subprocess.TimeoutExpired:
                            rospy.logwarn(f"进程组 PGID {pgid} (PID {pid}) 在 SIGTERM 后仍未终止。强制发送 SIGKILL...")
                            # 最后手段：强制杀死
                            os.killpg(pgid, signal.SIGKILL)
                            rospy.loginfo(f"已发送 SIGKILL 到 PGID {pgid}。")
                            process.wait() # 等待 SIGKILL 生效

                except ProcessLookupError:
                    # 进程在检查 poll() 和尝试 kill 之间已经消失了
                    rospy.loginfo(f"进程 PID {pid} 在尝试终止前已经不存在了。")
                except Exception as e:
                    rospy.logerr(f"终止进程 PID {pid} (PGID {pgid if 'pgid' in locals() else '未知'}) 时发生错误: {e}")
            else:
                # 如果进程在进入 shutdown_hook 时就已经结束了
                 rospy.loginfo(f"进程 PID {process.pid} 在关闭时已经结束。")


        self.launched_processes = [] # 清空列表
        rospy.loginfo(f"服务运行器 '{self.target_service_name}' 子进程清理完成。")

    def run(self):
        rospy.loginfo(f"服务 '{self.target_service_name}' 的运行器已激活，等待调用。")
        rospy.spin()
        # rospy.spin() 结束后，注册的 shutdown_hook 会被自动调用
        # 不需要在这里显式调用 self.shutdown_hook()

# --- 主程序入口部分保持不变 ---
if __name__ == '__main__':
    runner = None
    try:
        runner = SingleServiceRunner()
        runner.run()
    except rospy.ROSInterruptException:
        if runner:
             rospy.loginfo(f"接收到 {runner.target_service_name} 运行器的 ROS 中断信号。")
        else:
             rospy.loginfo("在初始化期间接收到 ROS 中断信号。")
    except SystemExit:
         rospy.logwarn("运行器因初始化期间的配置错误而退出。")
    except Exception as e:
        service_name = runner.target_service_name if runner and hasattr(runner, 'target_service_name') else "未知服务"
        try:
            rospy.logfatal(f"SingleServiceRunner 在处理 '{service_name}' 时发生未处理的异常: {e}")
        except:
            print(f"SingleServiceRunner 在处理 '{service_name}' 时发生严重错误: {e}")